import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

class KMeans:
    def __init__(self,k,upmit=1000):
        self.k=k
        self.centroids=None
        self.cluster_assment=None
        self.upmit=upmit

    def euclidean_distance(self,vecA, vecB):
        '''计算vecA与vecB之间的欧式距离'''
        # return np.sqrt(np.sum(np.square(vecA - vecB)))
        return np.linalg.norm(vecA - vecB)


    def random_centroids(self,data, k):
        ''' 随机创建k个中心点'''
        inex=np.random.choice(data.shape[0],k)
        return data[inex]


    def fit(self,data,distance_func=euclidean_distance):
        '''根据k-means算法求解聚类的中心'''
        num = np.shape(data)[0]  # 获得行数m
        feature_num=data.shape[1]
        self.cluster_assment = np.zeros((num, 1))  # 初试化一个矩阵，用来记录簇索引和存储距离平方
        self.centroids = self.random_centroids(data, k)  # 生成初始化点
        cluster_changed = True  # 判断是否需要重新计算聚类中心
        epoch=0
        while cluster_changed and epoch<self.upmit:
            epoch+=1
            print("%d/%d"%(self.upmit,epoch))
            distance = []
            for i in range(num):
                index_min = -1  # 所属的类别
                distance.append([self.euclidean_distance(data[i][0:feature_num],self.centroids[j]) for j in range(self.k)])
            self.cluster_assment=[np.where(distance[i]==min(distance[i]))[0][0] for i in range(len(distance))]
            data=np.hstack([data,np.array(self.cluster_assment).reshape([len(self.cluster_assment),1])])
            df=pd.DataFrame(data)
            for kClass, group in df.groupby(df.iloc[:,-1]):  # 根据某一列的值来聚类划分数据集，返回index
                self.centroids[int(kClass)]=np.array(group.iloc[:,0:feature_num]).mean(axis=0)

            if(data.shape[1]==feature_num+2):
                cluster_changed=(not np.all(data[:,-2]==data[:,-1]))
                data=np.hstack([data[:,0:feature_num],np.reshape(data[:,-1],[len(data),1])])
        return self.centroids, self.cluster_assment

def show_cluster(data, Y_train,k, centroids, cluster_assment):
    num, dim = data.shape
    mark = ['or', 'ob', 'og', 'oy', 'oc', 'om']
    df=pd.DataFrame(np.hstack([data,Y_train.reshape([len(Y_train),1])]))
    real_centroids=np.zeros([len(set(Y_train)),dim])
    for kClass, group in df.groupby(df.iloc[:, -1]):  # 根据某一列的值来聚类划分数据集，返回index
        real_centroids[int(kClass)] = np.array(group.iloc[:, 0:-1]).mean(axis=0)
        plt.plot(centroids[int(kClass)][0], centroids[int(kClass)][1], 'o', markeredgecolor='k', markersize=16)
    for i in range(num):
        mark_index = int(cluster_assment[i])
        plt.plot(data[i, 0], data[i, 1], mark[mark_index])
    for i in range(k):
        plt.plot(centroids[i, 0], centroids[i, 1], 'o', markeredgecolor='k', markersize=8)
    # print(real_centroids)
    # print(centroids)
    plt.show()

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
if __name__ == "__main__":
    iris = load_iris()
    X=np.array(iris.data)
    Y=iris.target
    X_train,X_test,Y_train,Y_test=train_test_split(X,Y,test_size=0.15)
    k=3
    km=KMeans(k)
    centroids, cluster_assment = km.fit(X_train)
    show_cluster(X_train,Y_train,k, centroids, cluster_assment)